#pragma once
#include<assert.h>

template<class K,class V>

struct AVLTreeNode
{
	AVLTreeNode<K, V>* _left;
	AVLTreeNode<K, V>* _right;
	AVLTreeNode<K, V>* _parent;              
	pair<K, V> _kv;                          // 存储的键对
	int _bf;                                 // balance factor

	AVLTreeNode(const pair<K, V>& kv)
		:_left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _kv(kv)
		, _bf(0)
	{}
};

template<class K, class V>
class AVLTree
{
	typedef AVLTreeNode<K, V> Node;
public:
	bool Insert(const pair<K, V>& kv);
	bool IsBalance();
	void InOrder();
	void Height();
	void Size();

private:
	void RotateL(Node* parent);              //左旋转
	void RotateR(Node* parent);              //右旋转
	void RotateRL(Node* parent);             //右左旋转
	void RotateLR(Node* parent);             //左右旋转
	bool _IsBalance(Node* root);
	void _InOrder(Node* root);
	int _Height(Node* root);
	int _Size(Node* root);


	Node* _root = nullptr;
};

template<class K, class V>
bool AVLTree<K, V>::Insert(const pair<K, V>& kv)
{
	if (_root == nullptr)
	{
		_root = new Node(kv);
		return true;
	}

	Node* parent = nullptr;
	Node* cur = _root;
	while (cur)
	{
		if (cur->_kv.first < kv.first)
		{
			parent = cur;
			cur = cur->_right;
		}
		else if (cur->_kv.first > kv.first)
		{
			parent = cur;
			cur = cur->_left;
		}
		else
		{
			return false;
		}
	}

	cur = new Node(kv);
	if (parent->_kv.first > kv.first)
	{
		parent->_left = cur;
	}
	else
	{
		parent->_right = cur;
	}
	cur->_parent = parent;

	// 更新平衡因子
	while (parent)
	{
		if (cur == parent->_right)
		{
			parent->_bf++;
		}
		else
		{
			parent->_bf--;
		}

		if (parent->_bf == 1 || parent->_bf == -1)
		{
			// 继续更新
			parent = parent->_parent;
			cur = cur->_parent;
		}
		else if (parent->_bf == 0)
		{
			break;
		}
		else if (parent->_bf == 2 || parent->_bf == -2)
		{
			// 需要旋转处理 -- 1、让这颗子树平衡 2、降低这颗子树的高度
			if (parent->_bf == 2 && cur->_bf == 1)
			{
				RotateL(parent);
			}
			else if (parent->_bf == -2 && cur->_bf == -1)
			{
				RotateR(parent);
			}
			else if (parent->_bf == -2 && cur->_bf == 1)
			{
				RotateLR(parent);
			}
			else if (parent->_bf == 2 && cur->_bf == -1)
			{
				RotateRL(parent);
			}
			else
			{
				assert(false);
			}

			break;
		}
		else
		{
			assert(false);
		}
	}

	return true;
}

template<class K, class V>
void AVLTree<K, V>::RotateL(Node* parent)
{
	Node* subR = parent->_right;
	Node* subRL = subR->_left;

	parent->_right = subRL;
	subR->_left = parent;

	Node* parentParent = parent->_parent;

	parent->_parent = subR;
	if (subRL)
		subRL->_parent = parent;

	if (_root == parent)
	{
		_root = subR;
		subR->_parent = nullptr;
	}
	else
	{
		if (parentParent->_left == parent)
		{
			parentParent->_left = subR;
		}
		else
		{
			parentParent->_right = subR;
		}

		subR->_parent = parentParent;
	}

	parent->_bf = subR->_bf = 0;
}

template<class K, class V>
void AVLTree<K, V>::RotateR(Node* parent)
{
	Node* subL = parent->_left;
	Node* subLR = subL->_right;

	parent->_left = subLR;
	if (subLR)
		subLR->_parent = parent;

	Node* parentParent = parent->_parent;

	subL->_right = parent;
	parent->_parent = subL;

	if (_root == parent)
	{
		_root = subL;
		subL->_parent = nullptr;
	}
	else
	{
		if (parentParent->_left == parent)
		{
			parentParent->_left = subL;
		}
		else
		{
			parentParent->_right = subL;
		}

		subL->_parent = parentParent;
	}

	subL->_bf = parent->_bf = 0;
}

template<class K, class V>
void AVLTree<K, V>::RotateRL(Node* parent)
{
	Node* subR = parent->_right;
	Node* subRL = subR->_left;
	int bf = subRL->_bf;

	RotateR(parent->_right);
	RotateL(parent);

	if (bf == 0)
	{
		// subRL自己就是新增
		parent->_bf = subR->_bf = subRL->_bf = 0;
	}
	else if (bf == -1)
	{
		// subRL的左子树新增
		parent->_bf = 0;
		subRL->_bf = 0;
		subR->_bf = 1;
	}
	else if (bf == 1)
	{
		// subRL的右子树新增
		parent->_bf = -1;
		subRL->_bf = 0;
		subR->_bf = 0;
	}
	else
	{
		assert(false);
	}
}

template<class K, class V>
void AVLTree<K, V>::RotateLR(Node* parent)
{
	Node* subL = parent->_left;
	Node* subLR = subL->_right;
	int bf = subLR->_bf;

	RotateL(parent->_left);
	RotateR(parent);

	if (bf == 0)
	{
		parent->_bf = subL->_bf = subLR->_bf = 0;
	}
	else if (bf == -1)
	{
		parent->_bf = 1;
		subLR->_bf = 0;
		subL->_bf = 0;
	}
	else if (bf == 1)
	{
		parent->_bf = 0;
		subLR->_bf = 0;
		subL->_bf = -1;
	}
	else
	{
		assert(false);
	}
}

template<class K, class V>
void AVLTree<K, V>::InOrder()
{
	_InOrder(_root);
	cout << endl;
}

template<class K, class V>
void AVLTree<K, V>::_InOrder(Node* root)
{
	if (root == nullptr)
		return;

	_InOrder(root->_left);
	cout << root->_kv.first << " ";
	_InOrder(root->_right);
}

template<class K, class V>
bool AVLTree<K, V>::IsBalance()
{
	return _IsBalance(_root);
}

template<class K, class V>
int AVLTree<K, V>::_Height(Node* root)
{
	if (root == nullptr)
		return 0;

	int leftHeight = _Height(root->_left);
	int rightHeight = _Height(root->_right);

	return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
}

template<class K, class V>
void AVLTree<K, V>::Height()
{
	cout << _Height(_root) << endl;
}

template<class K, class V>
void AVLTree<K, V>::Size()
{
	cout << _Size(_root) << endl;
}

template<class K, class V>
int AVLTree<K, V>::_Size(Node* root)
{
	return root == nullptr ? 0 :
		_Size(root->_left) +
		_Size(root->_right) + 1;
}
template<class K, class V>
bool AVLTree<K,V>::_IsBalance(Node* root)
{
	if (root == nullptr)
		return true;

	int leftHeight = _Height(root->_left);
	int rightHeight = _Height(root->_right);
	if (rightHeight - leftHeight != root->_bf)
	{
		cout << root->_kv.first << "平衡因子异常" << endl;
		return false;
	}

	return abs(rightHeight - leftHeight) < 2
		&& _IsBalance(root->_left)
		&& _IsBalance(root->_right);
}